#include "server.h"
#include <stdlib.h>
#include <string.h>

//===============================================================
// Starts the server on the specified port.
//===============================================================
CServer::CServer(int nPort) : m_nHighestFd(0)
{
	m_pciList = (CLIENTINFO **)malloc(MAX_CLIENTS * sizeof(CLIENTINFO *));
	if (!m_pciList)
	{
		// TODO: Shutdown, failed to allocate client list.
	}
	memset(m_pciList, 0, MAX_CLIENTS * sizeof(CLIENTINFO *));
	FD_ZERO(&m_fdSet);

	WSAStartup(MAKEWORD(2, 2), &m_wsaData);
	if (m_wsaData.wVersion != MAKEWORD(2, 2))
	{
		// TODO: Shutdown, bad WinSock version.
		return;
	}

	m_nSock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (m_nSock == INVALID_SOCKET)
	{
		// TODO: Shutdown, failed to create server socket.
		return;
	}

	memset(&m_saAddr, 0, sizeof(SOCKADDR_IN));
	m_saAddr.sin_family = AF_INET;
	m_saAddr.sin_port = htons(nPort);
	m_saAddr.sin_addr.s_addr = htonl(INADDR_ANY);

	if (bind(m_nSock, (SOCKADDR *)&m_saAddr, sizeof(m_saAddr)) == SOCKET_ERROR)
	{
		// TODO: Shutdown, failed to bind socket.
		return;
	}

	unsigned long nonBlocking = 1;
	ioctlsocket(m_nSock, FIONBIO, &nonBlocking);

	listen(m_nSock, SOMAXCONN);
	printf("Server started on port %d.\n", nPort);
}

//===============================================================
// Closes the server.
//===============================================================
CServer::~CServer()
{
	if (m_nSock != INVALID_SOCKET)
	{
		closesocket(m_nSock);
	}
	WSACleanup();
}

//===============================================================
// Updates all active connections.
//===============================================================
void CServer::Update()
{
	Accept();

    struct timeval tm;
    tm.tv_sec = tm.tv_usec = 0;

	memcpy(&m_fdRead, &m_fdSet, sizeof(fd_set));
	memcpy(&m_fdWrite, &m_fdSet, sizeof(fd_set));

	select(m_nHighestFd + 1, &m_fdRead, &m_fdWrite, NULL, &tm);
	for (int i = 0; i < MAX_CLIENTS; i++) 
	{
		if (m_pciList[i] == NULL)
			continue;

		/* Receive incoming data. */
		if (FD_ISSET(m_pciList[i]->fd, &m_fdRead))
		{
			if (!Receive(m_pciList[i]))
				continue;
		}

		/* Send data if there is any in the buffer.. */
		if (FD_ISSET(m_pciList[i]->fd, &m_fdWrite) && m_pciList[i]->buf_out_pos > 0)
		{
			if (!Send(m_pciList[i]))
				continue;
		}

		m_pciList[i]->client->Update();
	}
}

//===============================================================
// Accepts a incoming connection.
//===============================================================
void CServer::Accept()
{
	sockaddr_in saAddr;
	int addrLen = sizeof(SOCKADDR_IN);

	SOCKET nSock = accept(m_nSock, (SOCKADDR *)&saAddr, &addrLen);
	if (nSock != INVALID_SOCKET)
		Register(nSock, &saAddr);
}

//===============================================================
// Registers a new connection.
//===============================================================
void CServer::Register(SOCKET nSock, const SOCKADDR_IN *pAddr)
{
	u_int i;
	for (i = 0; i < MAX_CLIENTS; i++)
		if (m_pciList[i] == NULL)
			break;

	if (i >= MAX_CLIENTS) // TODO: Tell client server is full.
		return;

	FD_SET(nSock, &m_fdSet);
	if (nSock > m_nHighestFd)
		m_nHighestFd = nSock;

	m_pciList[i] = new CLIENTINFO;
	memset(m_pciList[i], 0, sizeof(CLIENTINFO));

	m_pciList[i]->id = i;
	m_pciList[i]->fd = nSock;
	memcpy(&(m_pciList[i]->addr), pAddr, sizeof(SOCKADDR_IN));
	m_pciList[i]->client = new CClient(m_pciList[i]);
	printf("A client has connected.\n");
}

//===============================================================
// Receives data from the client.
//===============================================================
bool CServer::Receive(CLIENTINFO *pClient)
{
	if (!pClient)
		return false;

	int nRecv = recv(pClient->fd, pClient->buf_in + pClient->buf_in_pos, (BUFFER_SIZE - pClient->buf_in_pos), 0);
	if (nRecv == SOCKET_ERROR)
	{
		int nErr = WSAGetLastError();
		switch (nErr)
		{
			case WSAENETDOWN:
			case WSAENETRESET:
			case WSAENOTCONN:
			case WSAEHOSTUNREACH:
			case WSAECONNABORTED:
			case WSAECONNRESET:
			case WSAETIMEDOUT:
			case WSAESHUTDOWN:
				/* If one of the cases above occurs the connection has been permanently lost so remove the client. */
				Terminate(pClient);
				return false;

			default:
				return true;
		}
	}

	/* If we receive 0, the connection has been closed. */
	if (nRecv == 0)
	{
		Terminate(pClient);
		return false;
	}
	else
	{
		pClient->buf_in_pos += nRecv;
		printf("Received %d bytes.\n", nRecv);
	}
	return true;
}

//===============================================================
// Sends data to the client.
//===============================================================
bool CServer::Send(CLIENTINFO *pClient)
{
	if (!pClient || pClient->buf_out_pos == 0)
		return false;

	if (send(pClient->fd, pClient->buf_out, pClient->buf_out_pos, 0) == SOCKET_ERROR)
	{
		int nErr = WSAGetLastError();
		switch (nErr) 
		{
			case WSAENETDOWN:
            case WSAENETRESET:
            case WSAENOTCONN:
            case WSAEHOSTUNREACH:
            case WSAECONNABORTED:
            case WSAECONNRESET:
            case WSAETIMEDOUT:
            case WSAESHUTDOWN:
				/* If one of the cases above occurs the connection has been permanently lost so remove the client. */
				Terminate(pClient); 
				return false;
		}
	}

	printf("Send %d bytes to client.\n", pClient->buf_out_pos);
	pClient->buf_out_pos = 0;
	return true;
}

//===============================================================
// Terminates the connection with the specified client.
//===============================================================
void CServer::Terminate(CLIENTINFO *pClient)
{
	if (!pClient || pClient->id < 0 || pClient->id >= MAX_CLIENTS)
		return;

	SOCKET nFd = pClient->fd;
	FD_CLR(pClient->fd, &m_fdSet);
	closesocket(pClient->fd);

	int nId = pClient->id;
	if (m_pciList[nId] != NULL)
	{
		delete (m_pciList[nId]->client);
		delete (m_pciList[nId]);
	}
	m_pciList[nId] = NULL;
	printf("A client has disconnected.\n");

	/* Find the new highest fd. */
	if (nFd == m_nHighestFd) 
	{
		m_nHighestFd = 0;
		for (int i = 0; i < MAX_CLIENTS; i++)
			if (m_pciList[i] != NULL && m_pciList[i]->fd > m_nHighestFd)
				m_nHighestFd = m_pciList[i]->fd;
	}
}
